from torch.utils.data import Dataset
import tqdm
import torch
import random


class BERTDataset(Dataset):
    """
    BERTDataset类用于准备BERT模型所需的训练数据集。
    它继承自torch.utils.data.Dataset，并实现了定制化的数据加载和处理方法。
    """

    def __init__(
        self,
        corpus_path,
        vocab,
        seq_len,
        encoding="utf-8",
        corpus_lines=None,
        on_memory=False,
    ):
        """
        初始化BERTDataset类。     
        如果on_memory为True，则将语料库加载到内存中。否则，计算语料库的行数。
        参数:
            corpus_path (str): 语料库文件的路径。
            vocab (Vocab): 词汇表对象，用于将单词映射到索引。
            seq_len (int): 序列长度，BERT输入序列的最大长度。
            encoding (str, optional): 文件编码，默认为'utf-8'。
            corpus_lines (int, optional): 语料库行数，如果提供，则在初始化时不会计算。
            on_memory (bool, optional): 是否将整个语料库加载到内存中，默认为True。
        """
        self.vocab = vocab
        self.seq_len = seq_len

        self.on_memory = on_memory
        self.corpus_lines = corpus_lines
        self.corpus_path = corpus_path
        self.encoding = encoding

        with open(corpus_path, "r", encoding=encoding) as f:
            if self.corpus_lines is None and not on_memory:
                for _ in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines):
                    self.corpus_lines += 1

            if on_memory:
                self.lines = [
                    line[:-1].split("\t")
                    for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines)
                ]
                print('第一行的内容',self.lines[0])
                print("第一行的数量:",len(self.lines[0]))
                self.corpus_lines = len(self.lines)
        """
            第一行的内容 ['Cuba to Get Rid of Dollars After a Decade', " HAVANA (Reuters) Cubans rushed to change dollars into  local pesos on Tuesday as President Fidel Castro's communist  government prepared to pull the U.S. currency from circulation  more than a decade after it was legalized here."]
            第一行的数量: 2
        """

        if not on_memory:
            self.file = open(corpus_path, "r", encoding=encoding)
            self.random_file = open(corpus_path, "r", encoding=encoding)

            for _ in range(
                random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)
            ):
                self.random_file.__next__()

    def __len__(self):
        """
        返回数据集的大小，即语料库的行数。
        
        返回:
            int: 语料库的行数。
        """
        return self.corpus_lines

    def __getitem__(self, item):
        """
        根据索引获取数据项，包括BERT输入序列、标签、段标签和下一句预测标签。
        
        参数:
            item (int): 数据项的索引。
            
        返回:
            dict: 包含BERT输入序列、标签、段标签和下一句预测标签的字典。
        """
        t1, t2, is_next_label = self.random_sent(item)
        t1_random, t1_label = self.random_word(t1)
        t2_random, t2_label = self.random_word(t2)


        """
            一个句子的头部加cls,尾部加eos标志
            label前面加pad,后面加pad
        """
        # [CLS] tag = SOS tag, [SEP] tag = EOS tag
        t1 = [self.vocab.sos_index] + t1_random + [self.vocab.eos_index]
        t2 = t2_random + [self.vocab.eos_index]

        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]
        t2_label = t2_label + [self.vocab.pad_index]
        # segment_label表示当前是第一句话还是第二句话,position的一部分
        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[
            : self.seq_len
        ]
        bert_input = (t1 + t2)[: self.seq_len]
        bert_label = (t1_label + t2_label)[: self.seq_len]

        padding = [self.vocab.pad_index for _ in range(self.seq_len - len(bert_input))]
        bert_input.extend(padding), bert_label.extend(padding), segment_label.extend(
            padding
        )
        
        output = {
            "bert_input": bert_input,
            "bert_label": bert_label,
            "segment_label": segment_label,
            "is_next": is_next_label,
        }

        return {key: torch.tensor(value) for key, value in output.items()}

    def random_word(self, sentence):
        """
        对给定句子中的单词进行随机处理，用于生成BERT的输入和标签。
        
        参数:
            sentence (str): 输入的句子。
            
        返回:
            tuple: 处理后的单词列表和对应的标签列表。
        """
        tokens = sentence.split()
        output_label = []

        for i, token in enumerate(tokens):
            prob = random.random()
            if prob < 0.15:
                prob /= 0.15

                # 80% randomly change token to mask token
                if prob < 0.8:
                    tokens[i] = self.vocab.mask_index

                # 10% randomly change token to random token
                elif prob < 0.9:
                    tokens[i] = random.randrange(len(self.vocab))

                # 10% randomly change token to current token
                else:
                    tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)

                output_label.append(self.vocab.stoi.get(token, self.vocab.unk_index))

            else:
                # self.vocab.stoi:单词到index的映射
                # 查询不到返回unk的index
                tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)
                output_label.append(0)

        return tokens, output_label

    def random_sent(self, index):
        """
        随机决定是否交换下一句，用于训练BERT的下一句预测任务。
        
        参数:
            index (int): 数据项的索引。
            
        返回:
            tuple: 两个句子和一个标签，指示是否是下一句。
        """
        t1, t2 = self.get_corpus_line(index)

        # output_text, label(isNotNext:0, isNext:1)
        if random.random() > 0.5:
            return t1, t2, 1
        else:
            return t1, self.get_random_line(), 0

    def get_corpus_line(self, item):
        """
        根据索引获取语料库中的句子对。
        
        参数:
            item (int): 数据项的索引。
            
        返回:
            tuple: 一个句子对。
        """
        if self.on_memory:
            #返回一句话中的上半和下半
            return self.lines[item][0], self.lines[item][1]
        else:
            line = self.file.__next__()
            if line is None:
                self.file.close()
                self.file = open(self.corpus_path, "r", encoding=self.encoding)
                line = self.file.__next__()

            t1, t2 = line[:-1].split("\t")
            return t1, t2

    def get_random_line(self):
        """
        获取语料库中的一个随机句子。
        
        返回:
            str: 一个随机句子。
        """
        if self.on_memory:
            return self.lines[random.randrange(len(self.lines))][1]

        line = self.file.__next__()
        if line is None:
            self.file.close()
            self.file = open(self.corpus_path, "r", encoding=self.encoding)
            for _ in range(
                random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)
            ):
                self.random_file.__next__()
            line = self.random_file.__next__()
        return line[:-1].split("\t")[1]